#!/bin/bash

PORT=7001
BASE_DIR="./received_files"
LOG_FILE="./file_transfer.log"
ZIP_DIR="$BASE_DIR/extracted_zips"
LIB_DIR="$BASE_DIR/libraries"
TEMP_DIR="$BASE_DIR/temp"

# 初始化目录
mkdir -p "$BASE_DIR" "$ZIP_DIR" "$LIB_DIR" "$TEMP_DIR"
touch "$LOG_FILE"

log_message() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] $message" | tee -a "$LOG_FILE"
}

# 检查端口是否可用
check_port() {
    if netstat -tuln 2>/dev/null | grep -q ":$PORT "; then
        log_message "错误: 端口 $PORT 已被占用"
        return 1
    fi
    return 0
}

# 检查必要工具
check_dependencies() {
    local deps=("file" "nc")
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            log_message "错误: 未找到 $dep 命令"
            return 1
        fi
    done
    
    if ! command -v unzip &> /dev/null; then
        log_message "警告: 未找到 unzip 命令，将无法处理 ZIP 文件"
    fi
    
    return 0
}

# 接收文件头部信息（文件名和文件大小）
receive_header() {
    local client_socket="$1"
    
    # 读取文件名长度（4字节）
    local filename_len_hex
    if ! read -r -n 8 filename_len_hex; then
        return 1
    fi
    
    # 将16进制转换为10进制
    local filename_len=$((16#$filename_len_hex))
    
    if [ "$filename_len" -eq 0 ] || [ "$filename_len" -gt 1024 ]; then
        log_message "错误: 无效的文件名长度: $filename_len"
        return 1
    fi
    
    # 读取文件名
    local filename
    if ! read -r -n "$filename_len" filename; then
        return 1
    fi
    
    # 读取文件大小（8字节）
    local filesize_hex
    if ! read -r -n 16 filesize_hex; then
        return 1
    fi
    
    local filesize=$((16#$filesize_hex))
    
    echo "$filename $filesize"
    return 0
}

# 接收文件内容
receive_file_content() {
    local client_socket="$1"
    local filename="$2"
    local filesize="$3"
    local output_path="$4"
    
    log_message "开始接收文件内容: $filename (大小: $filesize 字节)"
    
    # 使用dd接收指定大小的文件内容
    dd bs=4096 count=$((filesize / 4096 + 1)) of="$output_path" 2>/dev/null
    
    local actual_size=$(stat -c%s "$output_path" 2>/dev/null || echo 0)
    
    if [ "$actual_size" -eq "$filesize" ]; then
        log_message "文件内容接收成功: $filename (大小: $actual_size 字节)"
        return 0
    else
        log_message "错误: 文件大小不匹配 - 期望: $filesize, 实际: $actual_size"
        return 1
    fi
}

# 检测文件类型
detect_file_type() {
    local file="$1"
    local file_output=$(file -b "$file" 2>/dev/null)
    
    case "$file_output" in
        *"Zip archive"*)
            echo "zip"
            ;;
        *"ELF"*"shared object"*)
            echo "so"
            ;;
        *"ELF"*)
            echo "elf"
            ;;
        *)
            echo "unknown"
            ;;
    esac
}

# 处理 ZIP 文件
process_zip() {
    local file="$1"
    local filename="$2"
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local extract_dir="$ZIP_DIR/${filename%.*}_$timestamp"
    
    log_message "开始解压 ZIP 文件: $filename"
    
    if ! command -v unzip &> /dev/null; then
        log_message "错误: 未安装 unzip，无法解压文件"
        return 1
    fi
    
    mkdir -p "$extract_dir"
    
    if unzip -q "$file" -d "$extract_dir" 2>/dev/null; then
        local file_count=$(find "$extract_dir" -type f | wc -l)
        log_message "ZIP 文件解压成功: $extract_dir (包含 $file_count 个文件)"
        
        # 列出解压的文件
        find "$extract_dir" -type f | head -10 | while IFS= read -r extracted_file; do
            log_message "  - $(basename "$extracted_file")"
        done
        
        if [ "$file_count" -gt 10 ]; then
            log_message "  ... 还有 $((file_count - 10)) 个文件"
        fi
        
        return 0
    else
        log_message "错误: ZIP 文件解压失败: $filename"
        return 1
    fi
}

# 处理 SO 文件
process_so() {
    local file="$1"
    local filename="$2"
    local timestamp=$(date +%Y%m%d_%H%M%S)
    local so_name=$(basename "$file")
    local dest_file="$LIB_DIR/${so_name%.*}_${timestamp}.so"
    
    log_message "处理共享库文件: $filename"
    
    # 复制到库目录
    cp "$file" "$dest_file"
    
    # 设置适当的权限
    chmod 755 "$dest_file"
    
    # 获取文件信息
    local file_info=$(file "$dest_file")
    local file_size=$(stat -c%s "$dest_file")
    
    log_message "共享库已保存: $dest_file"
    log_message "文件信息: $file_info"
    log_message "文件大小: $file_size 字节"
    
    # 如果是 ELF 文件，尝试获取更多信息
    if command -v readelf >/dev/null 2>&1; then
        local elf_info=$(readelf -h "$dest_file" 2>/dev/null | grep -E "Type:|Machine:" | head -2)
        if [ -n "$elf_info" ]; then
            log_message "ELF 信息: $elf_info"
        fi
    fi
    
    return 0
}

# 处理未知文件类型
process_unknown() {
    local file="$1"
    local filename="$2"
    local file_type=$(file -b "$file")
    
    log_message "未知文件类型: $filename"
    log_message "检测到的类型: $file_type"
    
    # 移动到未知文件目录
    local unknown_dir="$BASE_DIR/unknown_types"
    mkdir -p "$unknown_dir"
    mv "$file" "$unknown_dir/$filename"
    
    log_message "文件已移动到: $unknown_dir/$filename"
}

# 处理接收到的文件
process_received_file() {
    local temp_file="$1"
    local original_filename="$2"
    local filesize="$3"
    
    local final_file="$BASE_DIR/$original_filename"
    
    # 移动临时文件到最终位置
    mv "$temp_file" "$final_file"
    
    log_message "文件保存为原始名称: $final_file"
    
    # 检测文件类型
    local file_type=$(detect_file_type "$final_file")
    log_message "检测到的文件类型: $file_type"
    
    # 根据文件类型处理
    case "$file_type" in
        "zip")
            process_zip "$final_file" "$original_filename"
            ;;
        "so"|"elf")
            process_so "$final_file" "$original_filename"
            ;;
        *)
            log_message "无需特殊处理的文件类型: $file_type"
            ;;
    esac
}

# 主接收循环
main_receive_loop() {
    while true; do
        log_message "等待客户端连接..."
        
        # 使用nc监听端口
        nc -l -p $PORT -c "
            echo '客户端连接成功'
            
            # 接收文件名长度（4字节，16进制）
            read -r -n 8 filename_len_hex
            filename_len=\$((16#\$filename_len_hex))
            echo \"文件名长度: \$filename_len\"
            
            # 接收文件名
            read -r -n \$filename_len filename
            echo \"文件名: \$filename\"
            
            # 接收文件大小（8字节，16进制）
            read -r -n 16 filesize_hex
            filesize=\$((16#\$filesize_hex))
            echo \"文件大小: \$filesize\"
            
            # 生成临时文件名
            temp_file=\"/tmp/receive_\$\$_\$(date +%s).dat\"
            
            # 接收文件内容
            dd bs=4096 count=\$((filesize / 4096 + 1)) of=\"\$temp_file\" 2>/dev/null
            
            # 发送处理结果回显
            actual_size=\$(stat -c%s \"\$temp_file\" 2>/dev/null || echo 0)
            if [ \"\$actual_size\" -eq \"\$filesize\" ]; then
                echo 'SUCCESS:文件接收成功'
            else
                echo 'ERROR:文件大小不匹配'
            fi
        " | while IFS= read -r line; do
            case "$line" in
                "客户端连接成功")
                    log_message "检测到新的客户端连接"
                    ;;
                "文件名长度:"*)
                    filename_len="${line#文件名长度: }"
                    ;;
                "文件名:"*)
                    filename="${line#文件名: }"
                    ;;
                "文件大小:"*)
                    filesize="${line#文件大小: }"
                    ;;
                "SUCCESS:"*)
                    log_message "${line#SUCCESS:}"
                    if [ -n "$filename" ] && [ -n "$filesize" ]; then
                        # 在实际应用中，这里需要处理接收到的文件
                        log_message "准备处理文件: $filename (大小: $filesize)"
                    fi
                    ;;
                "ERROR:"*)
                    log_message "${line#ERROR:}"
                    ;;
                *)
                    log_message "客户端消息: $line"
                    ;;
            esac
        done
        
        sleep 1
    done
}

# 简化版接收循环（兼容性更好）
simple_receive_loop() {
    while true; do
        log_message "等待客户端连接..."
        
        # 使用临时文件接收数据
        TEMP_FILE=$(mktemp -p "$TEMP_DIR")
        
        # 监听端口并接收数据
        if nc -l -p $PORT > "$TEMP_FILE" 2>/dev/null; then
            if [ -s "$TEMP_FILE" ]; then
                local file_size=$(stat -c%s "$TEMP_FILE")
                log_message "接收到数据流 (大小: $file_size 字节)"
                
                # 这里应该根据协议解析文件名
                # 由于nc限制，我们使用时间戳和检测到的类型
                local timestamp=$(date +%Y%m%d_%H%M%S)
                local file_type=$(detect_file_type "$TEMP_FILE")
                local final_filename="received_${timestamp}.${file_type}"
                local final_file="$BASE_DIR/$final_filename"
                
                mv "$TEMP_FILE" "$final_file"
                
                log_message "文件保存为: $final_filename (类型: $file_type)"
                
                # 根据文件类型处理
                case "$file_type" in
                    "zip")
                        process_zip "$final_file" "$final_filename"
                        ;;
                    "so"|"elf")
                        process_so "$final_file" "$final_filename"
                        ;;
                    *)
                        log_message "未知文件类型，已保存但未特殊处理"
                        ;;
                esac
            else
                rm -f "$TEMP_FILE"
                log_message "接收到空文件"
            fi
        else
            log_message "接收连接失败"
            sleep 1
        fi
    done
}

# 显示统计信息
show_stats() {
    log_message "=== 文件接收统计 ==="
    log_message "ZIP 文件数量: $(find "$ZIP_DIR" -type d | grep -v "^$ZIP_DIR$" | wc -l)"
    log_message "SO 文件数量: $(find "$LIB_DIR" -name "*.so" | wc -l)"
    log_message "总接收文件数: $(find "$BASE_DIR" -maxdepth 1 -type f | wc -l)"
}

# 清理临时文件
cleanup() {
    log_message "正在清理临时文件..."
    find "$TEMP_DIR" -name "*.dat" -delete
    log_message "清理完成"
}

# 主函数
main() {
    log_message "=== 启动增强版文件接收服务 ==="
    log_message "监听端口: $PORT"
    log_message "文件保存目录: $BASE_DIR"
    log_message "ZIP 解压目录: $ZIP_DIR"
    log_message "库文件目录: $LIB_DIR"
    
    if ! check_port; then
        exit 1
    fi
    
    if ! check_dependencies; then
        exit 1
    fi
    
    # 显示初始统计
    show_stats
    
    # 设置信号处理
    trap 'log_message "服务停止"; show_stats; cleanup; exit 0' INT TERM
    
    # 启动时清理临时文件
    cleanup
    
    log_message "选择接收模式:"
    log_message "1. 简单模式 (兼容性好)"
    log_message "2. 增强模式 (需要客户端配合)"
    
    # 默认使用简单模式
    # 如果您有能够发送文件名的客户端，可以取消注释下面的行来使用增强模式
    # read -p "请选择模式 (1/2): " mode
    
    mode=1  # 默认使用简单模式
    
    case "$mode" in
        2)
            log_message "启动增强接收模式"
            main_receive_loop
            ;;
        *)
            log_message "启动简单接收模式"
            simple_receive_loop
            ;;
    esac
}

# 运行主函数
main
